/*
 *        (C) COPYRIGHT Ingenic Limited.
 *             ALL RIGHTS RESERVED
 *
 * File       : inference_yolov5s.cpp
 * Authors    : ljhong, cywang
 * Create Time: 2024-11-21:15:57:21
 * Description:
 * 
 */

#include "ingenic_mnni.h"
#include "common/common_utils.h"
#include "utils/imgproc.h"
#include "utils/postproc.h"

#include <cmath>
#include <errno.h>
#include <fstream>
#include <iostream>
#include <memory>
#include <stdio.h>
#include <stdlib.h>
#include <string>
#include <cstring>
#include <string.h>
#include <vector>
#include <algorithm>
#include <sys/stat.h>   
#include <sys/types.h>
#include <dirent.h>     
#include <unistd.h>     
#include <sstream>      
#include <iomanip>      
#include <numeric>  // 添加 numeric 头文件，用于 std::inner_product
#include "yolo_interface.h"
#include "recog_interface.h"
#include "xb_image_utils.h"

// // #define IS_ALIGN_64(x) (((size_t)x) & 0x3F)
#ifdef URANUS_NV12
magik::uranus::ChannelLayout DATA_FORMAT = magik::uranus::ChannelLayout::NV12;
#else
magik::uranus::ChannelLayout DATA_FORMAT = magik::uranus::ChannelLayout::RGBA;
#endif


using namespace std;
using namespace magik::uranus;


typedef unsigned char uint8_t;
#define RET_CHECKER(ret, msg)                                           \
    do {                                                                \
        if (!(ret)) {                                                   \
            fprintf(stderr, "%s [%s %d]\n", msg, __func__, __LINE__);   \
            return -1;                                                  \
        }                                                               \
    } while(0)

int main(int argc, char **argv) {
    std::string face_det_model_path = argv[1];
    std::string recog_model_path = argv[2];
    std::string image_path = argv[3];
    if (argc != 4)
    {
        printf("%s face_det_model_path recog_model_path image \n", argv[0]);
        exit(0);
    }
    //  加载检测模型
    YOLOV5Config yolo5_config;
    yolo5_config.conf_thresh = 0.6f;       // 修改置信度阈值
    yolo5_config.nms_thres = 0.45f;        // 修改NMS阈值
    yolo5_config.classes = 1;              // 修改类别数
    yolo5_config.model_path = face_det_model_path; // 设置模型路径
    yolo5_config.input_h = 640;
    yolo5_config.input_w = 640;

    auto yolov5_interface = YOLOV5Interface(yolo5_config);
    yolov5_interface.loadModel();
    yolov5_interface.print_model_forward_memory();

    magik::uranus::shape_t working_img_shape = {1, yolo5_config.input_h, yolo5_config.input_w, 4};
    magik::uranus::Tensor working_img(working_img_shape);
    printf("load image: %s \n", image_path.c_str());
    FrameInfo input_info;
    MyImg img;
    printf("read_img=========\n");
    ImageUtils::read_img(input_info, img, image_path.c_str());

    std::vector<magik::uranus::ObjBbox_t> output_boxes;
    int ret = yolov5_interface.detect(input_info, output_boxes, DATA_FORMAT);


    //  加载识别模型
    RecogConfig recog_config;
    recog_config.dist_thres = 0.5;
    recog_config.model_path = recog_model_path;     
    recog_config.input_w = 224;    
    recog_config.input_h = 224; 
    recog_config.size = 128;

    auto recog_interface = RecogInterface(recog_config);
    recog_interface.loadModel();
    recog_interface.print_model_forward_memory();

    magik::uranus::Box_t tmp_box;
    tmp_box = {output_boxes[0].box.x0, output_boxes[0].box.y0, output_boxes[0].box.x1, output_boxes[0].box.y1};
    magik::uranus::shape_t recog_working_img_shape = {1, recog_config.input_h, recog_config.input_w, 4};
    magik::uranus::Tensor recog_working_img(recog_working_img_shape);
    YOLOUtils::prepare_data(input_info,  recog_working_img, tmp_box, recog_config.input_w, recog_config.input_h, DATA_FORMAT);
    
    MyImg img2;
    ImageUtils::frame_to_img(input_info, img2);
    ImageUtils::saveImg(img2, "frame_info.jpg", 60);

    MyImg img3;
    ImageUtils::tensor_to_img(recog_working_img, img3);
    ImageUtils::saveImg(img3, "recog_working_img.jpg", 60);

    std::vector<float> feat(recog_config.size);
    recog_interface.run_model(recog_working_img, feat);  
    for(int i=0; i< 10; i++){
        printf("feat[%d]: %f \n", i, feat[i]);
    }

    ImageUtils::draw_boxes(img, output_boxes);
    ImageUtils::saveImg(img, "output.jpg", 60);
    if(input_info.data != nullptr)
        free(input_info.data);
        
    recog_interface.releaseModel();
    yolov5_interface.releaseModel();
    return 0;
}

